import os
import unittest

import torch
from torch import hub
from torch.backends import cudnn
from torchsummary import summary

from config import Config
from training.models.sifdnetv2 import ALAM, sifdnet_v2
from training.models.sifdnetv2_2 import GradientBoostBlock, sifdnet_v2_2

CONFIG = Config()
hub.set_dir(CONFIG['TORCH_HOME'])
os.environ["CUDA_VISIBLE_DEVICES"] = CONFIG['CUDA_VISIBLE_DEVICES']

torch.backends.cudnn.benchmark = True


class SifdNetTestCase(unittest.TestCase):

    def test_sifdnet_v2(self):
        self.assertTrue(torch.cuda.is_available())
        model = sifdnet_v2(pretrained=True)
        model = model.cuda()
        input_size = model.default_cfg['input_size']
        summary(model, input_size=input_size)

    def test_alam1(self):
        self.assertTrue(torch.cuda.is_available())
        model = ALAM(in_channels=448, out_channels=160, size=(16, 16))
        model = model.cuda()
        input_size = [(160, 16, 16), (448, 8, 8)]
        summary(model, input_size=input_size)

    def test_alam2(self):
        self.assertTrue(torch.cuda.is_available())
        model = ALAM(in_channels=32, out_channels=24, size=(128, 128))
        model = model.cuda()
        input_size = [(24, 128, 128), (32, 64, 64)]
        summary(model, input_size=input_size)

    def test_gradient_boost_block(self):
        self.assertTrue(torch.cuda.is_available())
        model = GradientBoostBlock(in_channels=24, out_channels=32, pwl_stride=2, pwl_padding=1)
        model = model.cuda()
        input_size = (24, 128, 128)
        summary(model, input_size=input_size)

    def test_sifdnet_v2_2(self):
        self.assertTrue(torch.cuda.is_available())
        model = sifdnet_v2_2(pretrained=True)
        model = model.cuda()
        input_size = model.default_cfg['input_size']
        summary(model, input_size=input_size)

    def test_print_require_grad(self):
        self.assertTrue(torch.cuda.is_available())
        model = sifdnet_v2_2(pretrained=True)
        model.cuda()
        params = [param for param in model.sobel_stream.parameters()]
        print(len(params))


if __name__ == '__main__':
    unittest.main()
